import os
import time
import datetime
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import requests
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import matplotlib.pyplot as plt

# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 定义 LSTM 模型
class LSTMNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(LSTMNet, self).__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        output = self.fc(lstm_out[:, -1, :])  # 取最后一个时间步的输出
        return output

    def preprocess_data(self, data):
        """
        预处理股票数据，生成特征和目标值
        :param data: 包含股票数据的列表
        :return: 特征张量和目标张量
        """
        features = []
        targets = []
        for i in range(len(data) - 5):  # 确保有足够的数据来形成6天的序列
            input_data = []
            for j in range(i, i + 5):  # 前5天作为特征
                day_data = data[j]
                input_data.append([
                    float(day_data['open']),
                    float(day_data['high']),
                    float(day_data['low']),
                    float(day_data['close']),
                    float(day_data['volume'])
                ])
            features.append(input_data)
            # 第6天的收盘价作为目标
            target_data = data[i + 5]
            targets.append(float(target_data['close']))

        # 转换为张量
        features = torch.tensor(features, dtype=torch.float32)  # 形状为 (样本数, 5, 5)
        targets = torch.tensor(targets, dtype=torch.float32).view(-1, 1)  # 目标值为1维
        return features, targets

    def train_model(self, train_loader, num_epochs=10, learning_rate=0.001):
        """
        训练模型的方法
        :param train_loader: DataLoader，训练数据加载器
        :param num_epochs: int，训练轮数
        :param learning_rate: float，学习率
        """
        # 定义损失函数和优化器
        criterion = nn.MSELoss()  # 回归问题使用均方误差损失
        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)

        self.train()  # 设置模型为训练模式
        for epoch in range(num_epochs):
            running_loss = 0.0
            for inputs, targets in train_loader:
                inputs, targets = inputs.to(device), targets.to(device)
                # 前向传播
                outputs = self(inputs)
                loss = criterion(outputs, targets)

                # 反向传播和优化
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item() * inputs.size(0)

            epoch_loss = running_loss / len(train_loader.dataset)
            print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss:.4f}")

    def predict(self, input_data):
        """
        预测方法
        :param input_data: 包含五天数据的列表，每个元素是一个字典
        :return: 预测结果（第6天的收盘价）
        """
        # 将输入数据转换为张量
        input_features = []
        for day_data in input_data:
            input_features.append([
                float(day_data['open']),
                float(day_data['high']),
                float(day_data['low']),
                float(day_data['close']),
                float(day_data['volume'])
            ])
        input_features = torch.tensor([input_features], dtype=torch.float32)  # 形状为 (1, 5, 5)

        # 设置模型为评估模式
        self.eval()
        with torch.no_grad():  # 关闭梯度计算
            prediction = self(input_features.to(device))
        return prediction.item()  # 返回预测值

# 获取股票数据
def getStockData():
    url = "https://api.mairui.club/hslt/list/b997d4403688d5e66a"
    response = requests.get(url)
    data = response.json()
    return data

# 获取历史数据
def getHistoryData(code):
    url = "http://money.finance.sina.com.cn/quotes_service/api/json_v2.php/CN_MarketData.getKLineData"
    params = {
        "symbol": code,
        "scale": "240",  # 时间周期，如5分钟、日线等
        "ma": "5",
        "datalen": "1024"
    }
    response = requests.get(url, params=params)
    if response.status_code == 200:
        return response.json()
    else:
        print(f"Error fetching data for {code}, status code: {response.status_code}")
        return None

# 主函数
def main():
    # 创建模型实例
    model = LSTMNet(input_dim=5, hidden_dim=128, output_dim=1).to(device)

    # 加载模型权重
    if os.path.exists("model_weights.pth"):
        state_dict = torch.load("model_weights.pth")
        model.load_state_dict(state_dict)
        print("Model weights loaded.")

    # 训练模型
    def dotrain():
        stocks = getStockData()
        stocks = stocks[0:30]  # 选择部分股票进行训练
        total = len(stocks)
        for index, stock in enumerate(stocks):
            code = f"{stock['jys']}{stock['dm']}"
            data = getHistoryData(code)
            if data is None:
                continue

            # 预处理数据
            features, targets = model.preprocess_data(data)
            dataset = TensorDataset(features, targets)
            train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

            # 训练模型
            model.train_model(train_loader, num_epochs=400, learning_rate=0.001)
            print(f"训练进度: {index / total * 100:.2f}%")

        # 保存模型权重
        torch.save(model.state_dict(), "model_weights.pth")
        print("Model weights saved.")


    # 预测
    def dopredict():
        code = "sh601933"
        # code = "sh605100"
        data = getHistoryData(code)
        if data is None:
            print("Error: No data available for prediction.")
            return
        offset = 100
        # 选择最近5天的数据作为输入
        input_data = data[-6 - offset:-1 -offset]
        predicted_price = model.predict(input_data)
        print(f"预测 {input_data[-1]['day']} 的收盘价为: {predicted_price:.2f}")

    # 执行训练和预测
    # dotrain()
    dopredict()

if __name__ == "__main__":
    main()